import pandas as pd
import numpy as np
from collections import deque
from datetime import datetime
import time as _time
from xtquant import xtdata


class THSNineTurn5m:
    """
    同花顺风格的九转策略 - 5分钟K线版本

    参数:
    - compare_offset: int, 比较偏移量(默认4，表示比较当前K线与4根前的K线)
    - required_bars: int, 触发信号所需的连续次数(默认9)
    - max_window_minutes: int, 最大保留数据窗口(分钟)
    - verbose: bool, 是否打印详细日志
    - restrict_session: bool, 是否限制交易时段(True则只在09:30-11:30和13:00-15:00计算)
    - enable_ma_filter: bool, 是否启用5分钟K线均价过滤(默认True)
    - ma_period: int, 5分钟K线周期(默认20根)
    - use_5min_ma: bool, 是否使用5分钟K线计算均价(默认True)
    """

    def __init__(self, compare_offset: int = 4, required_bars: int = 9,
                 max_window_minutes: int = 240, verbose: bool = True,
                 restrict_session: bool = True, enable_ma_filter: bool = False, ma_period: int = 60,
                 use_5min_ma: bool = True):
        self.compare_offset = int(compare_offset)
        self.required_bars = int(required_bars)
        self.max_window_minutes = int(max_window_minutes)
        self.verbose = bool(verbose)
        self.restrict_session = bool(restrict_session)
        self.enable_ma_filter = bool(enable_ma_filter)
        self.ma_period = int(ma_period)  # 5分钟K线周期，默认20根
        self.use_5min_ma = bool(use_5min_ma)  # 是否使用5分钟K线计算均价
        
        # 5分钟K线时间间隔（秒）
        self.bar_interval = 300  # 5分钟 = 300秒

        # 状态存储
        self._bar_state = {}  # code -> {bucket:int, open:float, high:float, low:float, close:float}
        self._closes = {}  # code -> deque of close prices
        self._counts = {}  # code -> {"up":int, "down":int}
        self.last_signal = {}  # code -> last signal type
        
        # 5分钟K线均价相关
        self._ma_cache = {}  # code -> {ma: float, last_update: float, cache_ttl: 300}
        self._ma_cache_ttl = 300  # 缓存5分钟

    def reset(self, code: str = None):
        """重置指定股票或全部股票的状态"""
        if code is None:
            self._bar_state.clear()
            self._closes.clear()
            self._counts.clear()
            self.last_signal.clear()
            self._ma_cache.clear()
        else:
            self._bar_state.pop(code, None)
            self._closes.pop(code, None)
            self._counts.pop(code, None)
            self.last_signal.pop(code, None)
            self._ma_cache.pop(code, None)

    def _log(self, code: str, message: str):
        """日志记录"""
        if self.verbose:
            print(f"【5分钟九转】{code} {message}")

    def _is_trading_time(self, dt: datetime) -> bool:
        """检查是否在交易时段内"""
        if not self.restrict_session:
            return True

        tm = dt.time()
        # A股交易时间: 09:30-11:30, 13:00-15:00
        in_morning = (datetime.strptime("09:30:00", "%H:%M:%S").time() <= tm <= datetime.strptime("11:30:00", "%H:%M:%S").time())
        in_afternoon = (datetime.strptime("13:00:00", "%H:%M:%S").time() <= tm <= datetime.strptime("15:00:00", "%H:%M:%S").time())
        return in_morning or in_afternoon

    def _get_ma(self, code: str, current_price: float) -> float:
        """
        获取5分钟K线均价
        
        参数:
        - code: 股票代码
        - current_price: 当前价格，用于缓存验证
        
        返回:
        - float: 5分钟K线均价，如果无法计算则返回None
        """
        if not self.enable_ma_filter:
            return 0.0  # 禁用过滤时返回0，确保所有信号都能通过
            
        current_time = _time.time()
        cache = self._ma_cache.get(code)
        
        # 检查缓存是否有效
        if cache and (current_time - cache['last_update']) < self._ma_cache_ttl:
            return cache['ma']
        
        # 根据配置选择计算方法
        if self.use_5min_ma:
            return self._get_5min_ma(code, current_time)
        else:
            return self._get_xt_ma(code, current_time)
    
    def _get_5min_ma(self, code: str, current_time: float) -> float:
        """使用5分钟K线计算均价"""
        try:
            # 获取最近N根5分钟K线数据
            end_date = datetime.now()
            start_date = end_date - pd.Timedelta(days=2)  # 取最近2天的数据
            
            # 下载5分钟K线数据
            xtdata.download_history_data(
                stock_code=code, 
                period="5m", 
                start_time=start_date.strftime("%Y%m%d"), 
                end_time=end_date.strftime("%Y%m%d")
            )
            _time.sleep(0.5)  # 等待下载完成
            
            # 获取5分钟K线数据
            data = xtdata.get_local_data(
                stock_list=[code], 
                period='5m', 
                start_time=start_date.strftime("%Y%m%d"), 
                end_time=end_date.strftime("%Y%m%d")
            )
            
            if data and code in data and data[code] is not None and not data[code].empty:
                df = data[code]
                
                # 识别收盘价列
                close_col = next((c for c in ['close', 'Close', 'CLOSE', 'lastPrice'] if c in df.columns), None)
                if close_col:
                    closes = df[close_col].dropna()
                    if len(closes) >= self.ma_period:
                        # 计算最近N根5分钟K线的收盘价平均值
                        recent_closes = closes.tail(self.ma_period)
                        ma = recent_closes.mean()
                        
                        # 更新缓存
                        self._ma_cache[code] = {
                            'ma': float(ma),
                            'last_update': current_time,
                            'cache_ttl': self._ma_cache_ttl
                        }
                        
                        self._log(code, f"最近{self.ma_period}根5分钟K线均价: {ma:.4f}")
                        return float(ma)
                    else:
                        self._log(code, f"5分钟K线数据不足{self.ma_period}根，当前只有{len(closes)}根")
                else:
                    self._log(code, "无法识别5分钟K线收盘价列")
            else:
                self._log(code, "未获取到5分钟K线数据")
                
        except Exception as e:
            self._log(code, f"计算5分钟K线均价异常: {e}")
        
        return None
    
    def _get_xt_ma(self, code: str, current_time: float) -> float:
        """使用xtdata计算日线均线（备用方案）"""
        try:
            # 下载最近N个交易日的数据
            end_date = datetime.now()
            start_date = end_date - pd.Timedelta(days=self.ma_period + 10)  # 多取几天确保有足够数据
            
            # 下载日线数据
            xtdata.download_history_data(
                stock_code=code, 
                period="1d", 
                start_time=start_date.strftime("%Y%m%d"), 
                end_time=end_date.strftime("%Y%m%d")
            )
            _time.sleep(0.5)  # 等待下载完成
            
            # 获取日线数据
            data = xtdata.get_local_data(
                stock_list=[code], 
                period='1d', 
                start_time=start_date.strftime("%Y%m%d"), 
                end_time=end_date.strftime("%Y%m%d")
            )
            
            if data and code in data and data[code] is not None and not data[code].empty:
                df = data[code]
                
                # 识别收盘价列
                close_col = next((c for c in ['close', 'Close', 'CLOSE', 'lastPrice'] if c in df.columns), None)
                if close_col:
                    closes = df[close_col].dropna()
                    if len(closes) >= self.ma_period:
                        # 计算均线
                        ma = closes.tail(self.ma_period).mean()
                        
                        # 更新缓存
                        self._ma_cache[code] = {
                            'ma': float(ma),
                            'last_update': current_time,
                            'cache_ttl': self._ma_cache_ttl
                        }
                        
                        self._log(code, f"xtdata {self.ma_period}日均线: {ma:.4f}")
                        return float(ma)
                    else:
                        self._log(code, f"数据不足{self.ma_period}天，当前只有{len(closes)}天数据")
                else:
                    self._log(code, "无法识别收盘价列")
            else:
                self._log(code, "未获取到日线数据")
                
        except Exception as e:
            self._log(code, f"计算xtdata {self.ma_period}日均线异常: {e}")
        
        return None

    def _validate_signal_with_ma(self, code: str, signal: str, current_price: float) -> str:
        """
        使用5分钟K线均价验证信号
        
        参数:
        - code: 股票代码
        - signal: 原始信号("买入"或"卖出")
        - current_price: 当前价格
        
        返回:
        - str: 验证后的信号，如果验证失败则返回None
        """
        if not self.enable_ma_filter:
            return signal
            
        ma = self._get_ma(code, current_price)
        if ma is None:
            # 无法获取均价时，记录日志但允许信号通过
            self._log(code, f"无法获取{self.ma_period}根5分钟K线均价，信号{signal}直接通过")
            return signal
        
        # 验证信号
        if signal == "买入":
            if current_price > ma:
                self._log(code, f"买入信号确立: 当前价格{current_price:.4f} > {self.ma_period}根5分钟K线均价{ma:.4f}")
                return signal
            else:
                self._log(code, f"买入信号被过滤: 当前价格{current_price:.4f} <= {self.ma_period}根5分钟K线均价{ma:.4f}")
                return None
        elif signal == "卖出":
            if current_price > ma:
                self._log(code, f"卖出信号确立: 当前价格{current_price:.4f} > {self.ma_period}根5分钟K线均价{ma:.4f}")
                return signal
            else:
                self._log(code, f"卖出信号被过滤: 当前价格{current_price:.4f} <= {self.ma_period}根5分钟K线均价{ma:.4f}")
                return None
        
        return signal

    def on_tick(self, code: str, timestamp, price) -> str:
        """
        处理tick数据，返回信号（按5分钟K线聚合）

        参数:
        - code: 股票代码
        - timestamp: 时间戳(支持多种格式)
        - price: 当前价格

        返回:
        - "买入" / "卖出" / None
        """
        if price is None:
            return None

        # 转换时间戳为datetime对象
        try:
            if isinstance(timestamp, (int, float)):
                # 处理时间戳(秒或毫秒)
                ts_sec = timestamp if timestamp < 1e10 else timestamp / 1000
                dt = datetime.fromtimestamp(ts_sec)
            else:
                # 处理字符串时间
                dt = pd.to_datetime(timestamp)
        except Exception:
            return None

        # 检查交易时段
        if not self._is_trading_time(dt):
            return None

        # 按5分钟聚合K线
        bar_bucket = int(dt.timestamp() // self.bar_interval)
        bar = self._bar_state.get(code)

        # 初始化状态
        cnt = self._counts.setdefault(code, {"up": 0, "down": 0})
        closes = self._closes.setdefault(code, deque(maxlen=1000))

        signal = None

        # 新5分钟K线开始
        if bar is None or bar_bucket != bar["bucket"]:
            if bar is not None:
                # 完成上一根5分钟K线
                closes.append(bar["close"])

                # 计算九转信号
                raw_signal = self._check_signal(code, closes, cnt)
                
                # 如果触发了九转信号，使用5分钟K线均价进行过滤
                if raw_signal:
                    signal = self._validate_signal_with_ma(code, raw_signal, bar["close"])

            # 开始新K线
            self._bar_state[code] = {
                "bucket": bar_bucket,
                "open": float(price),
                "high": float(price),
                "low": float(price),
                "close": float(price)
            }
        else:
            # 更新当前5分钟K线
            bar["close"] = float(price)
            if price > bar["high"]:
                bar["high"] = float(price)
            if price < bar["low"]:
                bar["low"] = float(price)

        # 获取当前连续计数
        up_count = cnt.get("up", 0)
        down_count = cnt.get("down", 0)
        ##print(f"九转检测{code}:{signal} 上涨计数:{up_count} 下跌计数:{down_count}")

        return signal

    def _check_signal(self, code: str, closes: deque, cnt: dict) -> str:
        """检查是否满足信号条件"""
        if len(closes) <= self.compare_offset:
            return None

        cur = closes[-1]
        ref = closes[-1 - self.compare_offset]

        # 更新计数：与同花顺更一致的容忍度
        # - 满足 cur >= ref 则看涨计数 +1，否则清零
        # - 满足 cur <= ref 则看跌计数 +1，否则清零
        # 注意：相等会同时累加两个方向，允许重叠序列
        # if cur >= ref:
        #     cnt["up"] += 1
        # else:
        #     cnt["up"] = 0
        #
        # if cur <= ref:
        #     cnt["down"] += 1
        # else:
        #     cnt["down"] = 0

        if cur > ref:  # 严格上涨
            cnt["up"] += 1
            cnt["down"] = 0  # 重置下跌计数
        elif cur < ref:  # 严格下跌
            cnt["down"] += 1
            cnt["up"] = 0  # 重置上涨计数
        else:  # 相等时，两个计数都清零
            cnt["up"] = 0
            cnt["down"] = 0

        # 检查信号
        signal = None
        if cnt["up"] >= self.required_bars:
            signal = "卖出"
            # 仅重置对应方向，允许另一方向的序列延续
            cnt["up"] = 0
            self._log(code, f"触发卖出信号，连续{self.required_bars}次5分钟K线收盘价不低于前{self.compare_offset}根K线")

        elif cnt["down"] >= self.required_bars:
            signal = "买入"
            cnt["down"] = 0
            self._log(code, f"触发买入信号，连续{self.required_bars}次5分钟K线收盘价不高于前{self.compare_offset}根K线")

        if signal:
            self.last_signal[code] = signal

        return signal

    def warmup_from_history(self, code: str, history_data: pd.DataFrame):
        """
        使用历史数据预热策略（按5分钟K线）

        参数:
        - code: 股票代码
        - history_data: DataFrame包含历史数据，必须有datetime索引和price列
        """
        if not isinstance(history_data.index, pd.DatetimeIndex):
            raise ValueError("历史数据必须有DatetimeIndex")

        # 按5分钟重采样
        df = history_data.resample('5min').agg({
            'price': 'ohlc'
        })
        df.columns = df.columns.droplevel(0)

        # 回放历史数据
        for _, row in df.iterrows():
            if pd.isna(row['close']):
                continue
            self.on_tick(code, row.name, row['close'])

    def get_down_count(self, code: str) -> int:
        """
        获取指定股票的当前下跌计数次数
        
        参数:
        - code: 股票代码
        
        返回:
        - int: 当前下跌计数次数，如果股票代码不存在则返回0
        """
        if code in self._counts:
            return self._counts[code].get("down", 0)
        return 0

    def get_up_count(self, code: str) -> int:
        """
        获取指定股票的当前上涨计数次数
        
        参数:
        - code: 股票代码
        
        返回:
        - int: 当前上涨计数次数，如果股票代码不存在则返回0
        """
        if code in self._counts:
            return self._counts[code].get("up", 0)
        return 0

    def get_ma(self, code: str) -> float:
        """
        获取指定股票的5分钟K线均价
        
        参数:
        - code: 股票代码
        
        返回:
        - float: 5分钟K线均价，如果无法获取则返回None
        """
        return self._get_ma(code, 0.0)

    def set_ma_filter_enabled(self, enabled: bool):
        """
        设置是否启用5分钟K线均价过滤
        
        参数:
        - enabled: bool, 是否启用过滤
        """
        self.enable_ma_filter = bool(enabled)
        if not enabled:
            self._log("", f"{self.ma_period}根5分钟K线均价过滤已禁用")
    
    def set_ma_calculation_method(self, use_5min: bool):
        """
        设置均价计算方式
        
        参数:
        - use_5min: bool, True使用5分钟K线，False使用日线
        """
        old_method = "5分钟K线" if self.use_5min_ma else "日线"
        self.use_5min_ma = bool(use_5min)
        new_method = "5分钟K线" if self.use_5min_ma else "日线"
        
        # 清除缓存，因为计算方式改变了
        self._ma_cache.clear()
        
        self._log("", f"均价计算方式从 {old_method} 切换到 {new_method}")
    
    def get_ma_method_info(self) -> str:
        """获取当前均价计算方式信息"""
        method = "5分钟K线" if self.use_5min_ma else "日线"
        return f"{method} {self.ma_period}根均价"


# 示例用法
if __name__ == "__main__":
    # 初始化同花顺九转策略（启用5分钟K线均价过滤）
    nine_turn = THSNineTurn5m(
        compare_offset=4,
        required_bars=9,
        verbose=True,
        enable_ma_filter=False,  # 启用5分钟K线均价过滤
        ma_period=20,  # 设置5分钟K线周期为20根
        use_5min_ma=True  # 使用5分钟K线计算均价
    )

    # 股票代码
    code = "513120.SH"

    # 下载当日tick数据
    today = datetime.now().strftime("%Y%m%d")
    xtdata.download_history_data(stock_code=code, period="tick", start_time=today, end_time=today)
    _time.sleep(1.5)
    data = xtdata.get_local_data(stock_list=[code], period='tick', start_time=today, end_time=today)

    if not data or code not in data or data[code] is None or data[code].empty:
        print(f"未获取到当日 {code} 的tick数据")
    else:
        df = data[code]
        # 识别时间/价格列
        time_col = next((c for c in df.columns if 'time' in c.lower()), None)
        price_col = next((c for c in ['price', 'lastPrice', 'close', 'last', 'current'] if c in df.columns), None)

        if not time_col or not price_col:
            print(f"{code} 数据列缺失，无法识别时间或价格列")
        else:
            # 统一时间列
            if str(df[time_col].dtype) in ('int64', 'float64'):
                df['__time__'] = pd.to_datetime(df[time_col], unit='ms') + pd.Timedelta(hours=8)
            else:
                df['__time__'] = pd.to_datetime(df[time_col], format='%Y%m%d%H%M%S', errors='coerce')
            df = df[df['__time__'].notna()].sort_values('__time__')

            # 回放tick数据并收集信号
            signals = []
            for _, row in df.iterrows():
                sig = nine_turn.on_tick(code, row['__time__'], row[price_col])
                if sig:
                    signals.append({
                        'time': row['__time__'].strftime('%Y-%m-%d %H:%M:%S'),
                        'type': sig,
                        'price': float(row[price_col])
                    })

            # 输出结果
            if not signals:
                print(f"当日 {code} 未产生九转买/卖信号")
            else:
                print(f"当日 {code} 九转信号：")
                for s in signals:
                    print(f"  {s['time']}  {s['type']}  价格: {s['price']:.4f}")

                buy_count = sum(1 for s in signals if s['type'] == '买入')
                sell_count = sum(1 for s in signals if s['type'] == '卖出')
                print(f"九转买入次数: {buy_count}，九转卖出次数: {sell_count}，合计: {len(signals)}")

            # 查询当前计数状态
            print(f"\n=== 当前九转计数状态 ===")
            current_down_count = nine_turn.get_down_count(code)
            current_up_count = nine_turn.get_up_count(code)
            print(f"当前 {code} 的下跌计数: {current_down_count}")
            print(f"当前 {code} 的上涨计数: {current_up_count}")
            print(f"说明: 下跌计数达到{nine_turn.required_bars}次会触发买入信号，上涨计数达到{nine_turn.required_bars}次会触发卖出信号")

            # 测试重置功能（如果需要清空状态重新开始，可以使用）
            # print(f"\n=== 测试重置功能 ===")
            # print(f"重置前 {code} 的下跌计数: {nine_turn.get_down_count(code)}")
            # nine_turn.reset(code)  # 重置指定股票
            # print(f"重置后 {code} 的下跌计数: {nine_turn.get_down_count(code)}")
            # nine_turn.reset()  # 重置所有股票
            # print(f"全部重置后 {code} 的下跌计数: {nine_turn.get_down_count(code)}")
            
            # 测试5分钟K线均价功能
            print(f"\n=== 测试{nine_turn.ma_period}根5分钟K线均价功能 ===")
            ma_value = nine_turn.get_ma(code)
            if ma_value is not None:
                print(f"当前 {code} 的{nine_turn.get_ma_method_info()}: {ma_value:.4f}")
                print(f"{nine_turn.ma_period}根5分钟K线均价过滤状态: {'启用' if nine_turn.enable_ma_filter else '禁用'}")
                
                # 测试禁用均价过滤
                nine_turn.set_ma_filter_enabled(False)
                print(f"禁用均价过滤后: {'启用' if nine_turn.enable_ma_filter else '禁用'}")
                
                # 重新启用
                nine_turn.set_ma_filter_enabled(True)
                print(f"重新启用均价过滤后: {'启用' if nine_turn.enable_ma_filter else '禁用'}")
                
                # 测试切换均价计算方式
                print(f"\n=== 测试均价计算方式切换 ===")
                print(f"当前计算方式: {nine_turn.get_ma_method_info()}")
                
                # 切换到日线
                nine_turn.set_ma_calculation_method(False)
                ma_value_daily = nine_turn.get_ma(code)
                if ma_value_daily is not None:
                    print(f"日线均价: {ma_value_daily:.4f}")
                
                # 切换回5分钟K线
                nine_turn.set_ma_calculation_method(True)
                ma_value_5min = nine_turn.get_ma(code)
                if ma_value_5min is not None:
                    print(f"5分钟K线均价: {ma_value_5min:.4f}")
                
                # 比较两种方式
                if ma_value_daily is not None and ma_value_5min is not None:
                    diff = abs(ma_value_daily - ma_value_5min)
                    print(f"两种方式差异: {diff:.4f} ({diff/ma_value_daily*100:.2f}%)")
            else:
                print(f"无法获取 {code} 的{nine_turn.ma_period}根5分钟K线均价数据")